import os
import json
import time
from typing import List, Dict, Optional
from datetime import datetime
from pymilvus import connections, Collection, utility

from config.settings import settings
from app.core.logging import get_logger

logger = get_logger(__name__)

class KnowledgeBaseManager:
    """知识库管理器"""
    
    def __init__(self):
        self.kb_config_file = "data/knowledge_bases.json"
        self._ensure_config_file()
    
    def _ensure_config_file(self):
        """确保配置文件存在"""
        os.makedirs(os.path.dirname(self.kb_config_file), exist_ok=True)
        if not os.path.exists(self.kb_config_file):
            # 创建默认知识库配置
            default_kb = {
                "default": {
                    "name": "default",
                    "description": "默认知识库",
                    "created_time": datetime.now().isoformat(),
                    "last_updated": datetime.now().isoformat(),
                    "file_count": 0,
                    "text_collection": settings.MILVUS_COLLECTION,
                    "image_collection": settings.MILVUS_COLLECTION_IMAGE
                }
            }
            self._save_config(default_kb)
    
    def _load_config(self) -> Dict:
        """加载知识库配置"""
        try:
            with open(self.kb_config_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载知识库配置失败: {e}")
            return {}
    
    def _save_config(self, config: Dict):
        """保存知识库配置"""
        try:
            with open(self.kb_config_file, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存知识库配置失败: {e}")
    
    def create_knowledge_base(self, name: str, description: str = "") -> bool:
        """
        创建知识库
        
        Args:
            name: 知识库名称
            description: 知识库描述
            
        Returns:
            bool: 是否创建成功
        """
        try:
            config = self._load_config()
            
            if name in config:
                logger.warning(f"知识库 {name} 已存在")
                return False
            
            # 创建知识库配置
            kb_config = {
                "name": name,
                "description": description,
                "created_time": datetime.now().isoformat(),
                "last_updated": datetime.now().isoformat(),
                "file_count": 0,
                "text_collection": settings.MILVUS_COLLECTION,
                "image_collection": settings.MILVUS_COLLECTION_IMAGE
            }
            
            config[name] = kb_config
            self._save_config(config)
            
            # 创建Milvus分区
            self._create_milvus_partitions(name)
            
            logger.info(f"知识库 {name} 创建成功")
            return True
            
        except Exception as e:
            logger.error(f"创建知识库 {name} 失败: {e}")
            return False
    
    def _create_milvus_partitions(self, kb_name: str):
        """创建Milvus分区"""
        try:
            from pymilvus import Collection
            
            # 获取文本集合
            text_collection = Collection(settings.MILVUS_COLLECTION)
            
            # 创建文本分区
            if not text_collection.has_partition(kb_name):
                text_collection.create_partition(kb_name)
                logger.info(f"创建文本分区: {kb_name}")
            
            # 获取图片集合
            image_collection = Collection(settings.MILVUS_COLLECTION_IMAGE)
            
            # 创建图片分区
            if not image_collection.has_partition(kb_name):
                image_collection.create_partition(kb_name)
                logger.info(f"创建图片分区: {kb_name}")
            
            logger.info(f"知识库 {kb_name} 的Milvus分区创建成功")
            
        except Exception as e:
            logger.error(f"创建知识库 {kb_name} 的Milvus分区失败: {e}")
            raise
    
    def get_knowledge_base(self, name: str) -> Optional[Dict]:
        """
        获取知识库信息
        
        Args:
            name: 知识库名称
            
        Returns:
            Dict: 知识库信息
        """
        config = self._load_config()
        return config.get(name)
    
    def list_knowledge_bases(self) -> List[Dict]:
        """
        列出所有知识库
        
        Returns:
            List[Dict]: 知识库列表
        """
        config = self._load_config()
        return list(config.values())
    
    def delete_knowledge_base(self, name: str) -> bool:
        """
        删除知识库
        
        Args:
            name: 知识库名称
            
        Returns:
            bool: 是否删除成功
        """
        try:
            config = self._load_config()
            
            if name not in config:
                logger.warning(f"知识库 {name} 不存在")
                return False
            
            # 删除Milvus分区
            self._delete_milvus_partitions(name)
            
            # 删除配置
            del config[name]
            self._save_config(config)
            
            logger.info(f"知识库 {name} 删除成功")
            return True
            
        except Exception as e:
            logger.error(f"删除知识库 {name} 失败: {e}")
            return False
    
    def _delete_milvus_partitions(self, kb_name: str):
        """删除Milvus分区"""
        try:
            from pymilvus import Collection
            
            # 删除文本分区
            text_collection = Collection(settings.MILVUS_COLLECTION)
            if text_collection.has_partition(kb_name):
                text_collection.drop_partition(kb_name)
                logger.info(f"删除文本分区: {kb_name}")
            
            # 删除图片分区
            image_collection = Collection(settings.MILVUS_COLLECTION_IMAGE)
            if image_collection.has_partition(kb_name):
                image_collection.drop_partition(kb_name)
                logger.info(f"删除图片分区: {kb_name}")
                
        except Exception as e:
            logger.error(f"删除知识库 {kb_name} 的Milvus分区失败: {e}")
    
    def update_knowledge_base_stats(self, name: str, file_count: int = None):
        """
        更新知识库统计信息
        
        Args:
            name: 知识库名称
            file_count: 文件数量
        """
        try:
            config = self._load_config()
            
            if name not in config:
                logger.warning(f"知识库 {name} 不存在")
                return
            
            if file_count is not None:
                config[name]["file_count"] = file_count
            
            config[name]["last_updated"] = datetime.now().isoformat()
            self._save_config(config)
            
        except Exception as e:
            logger.error(f"更新知识库 {name} 统计信息失败: {e}")
    
    def get_collection_names(self, kb_name: str) -> Dict[str, str]:
        """
        获取知识库的集合名称（所有知识库使用相同的集合，通过分区区分）
        
        Args:
            kb_name: 知识库名称
            
        Returns:
            Dict: 集合名称映射
        """
        # 所有知识库都使用相同的集合名称
        return {
            "text_collection": settings.MILVUS_COLLECTION,
            "image_collection": settings.MILVUS_COLLECTION_IMAGE
        }

# 全局知识库管理器实例
_kb_manager = None

def get_kb_manager() -> KnowledgeBaseManager:
    """获取知识库管理器实例"""
    global _kb_manager
    if _kb_manager is None:
        _kb_manager = KnowledgeBaseManager()
    return _kb_manager 